"""
@File    : open_shp.py
@Time    : 2025/9/15 上午9:29
@Author  : xiashuobad
@Desc    : 

/**
 * _ooOoo_
 * o8888888o
 * 88" . "88
 * (| -_- |)
 *  O\ = /O
 * ___/`---'\____
 * .   ' \\| |// `.
 * / \\||| : |||// \
 * / _||||| -:- |||||- \
 * | | \\\ - /// | |
 * | \_| ''\---/'' | |
 * \ .-\__ `-` ___/-. /
 * ___`. .' /--.--\ `. . __
 * ."" '< `.___\_<|>_/___.' >'"".
 * | | : `- \`.;`\ _ /`;.`/ - ` : | |
 * \ \ `-. \_ __\ /__ _/ .-` / /
 * ======`-.____`-.___\_____/___.-`____.-'======
 * `=---=' bug泛滥 佛已瘫痪
"""
from pathlib import Path
from typing import Literal
from osgeo import ogr, osr, gdal
from typing import Generator, Iterable
from .utils import GeoConverter, require_attr_not_empty, copy_data_source


class OpenShp:
    open_mode = Literal['r', 'u', 'w']

    def __init__(self, fp: str, mode: open_mode = 'r'):
        self._fp = Path(fp)
        self._mode = self._checked_mode(mode)
        self._ds = None
        self._open()

    def _open(self):
        if self.mode != 'w':
            self._ds = gdal.VectorTranslate(
                '',
                self.file_path,
                format='Memory',
            )
            if self._ds is None:
                raise OSError(f"无法打开文件: {self.file_path}")

    def save(self, save_path=None):
        if self.mode == 'r':
            raise OSError("只读模式下无法保存数据！")
        if self.data_source is None:
            raise OSError("数据源为空！")
        save_path = save_path or self.file_path
        Path(save_path).parent.mkdir(parents=True, exist_ok=True)
        if isinstance(self.data_source, gdal.Dataset):
            gdal.VectorTranslate(
                save_path,  # 输出文件路径
                self.data_source,  # 内存数据集
                format="ESRI Shapefile",  # 输出格式
                layerCreationOptions=["ENCODING=UTF-8"]
            )
        elif isinstance(self.data_source, ogr.DataSource):
            copy_data_source(self.data_source, save_path)
        else:
            raise TypeError("数据源类型错误！")

    def _checked_mode(self, mode: open_mode) -> open_mode:
        if mode not in self.open_mode.__args__:
            raise OSError(f'不支持的打开模式：{mode}, 请使用以下模式：{self.open_mode.__args__}！')
        return mode

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    @property
    def mode(self) -> open_mode:
        return self._mode

    @property
    def file_path(self) -> str:
        return str(self._fp)

    @property
    def file_name(self) -> str:
        return self._fp.stem

    @property
    def data_source(self) -> ogr.DataSource:
        return self._ds

    @property
    def layer(self) -> ogr.Layer:
        """获取第一个图层"""
        if self._ds is not None:
            return self._ds.GetLayer(0)

    @property
    def extent(self) -> tuple[float, float, float, float]:
        """获取地理范围 (min_x, max_x, min_y, max_y)"""
        if self.layer is not None:
            return self.layer.GetExtent()

    @property
    def geo_bounds(self) -> tuple[float, float, float, float]:
        """返回 (left, top, right, bottom) 格式的地理范围"""
        extent = self.extent
        if extent is not None:
            min_x, max_x, min_y, max_y = extent
            return min_x, max_y, max_x, min_y  # left, top, right, bottom

    @property
    def spatial_ref(self) -> osr.SpatialReference:
        """获取当前空间参考系统"""
        if self.layer is not None:
            return self.layer.GetSpatialRef()

    @property
    def geo_projection(self) -> str:
        """获取当前地理投影系统"""
        if self.spatial_ref is not None:
            return self.spatial_ref.ExportToWkt()

    @property
    def features(self) -> Generator[ogr.Feature, None, None]:
        if self.layer is not None:
            for feature in self.layer:
                yield feature

    @property
    def feature_count(self) -> int:
        if self.layer is not None:
            return self.layer.GetFeatureCount()

    def set_data_source(self, data_source: 'OpenShp|gdal.Dataset|ogr.DataSource|str'):
        """
        设置数据源
        :param data_source: 支持OpenShp对象，gdal.Dataset对象，ogr.DataSource或者shp文件路径
        :return:
        """
        if isinstance(data_source, OpenShp):
            data_source = data_source.data_source
        if isinstance(data_source, (str, gdal.Dataset)):
            self._ds = gdal.VectorTranslate(
                '',
                data_source,
                format='Memory',
            )
        elif isinstance(data_source, ogr.DataSource):
            shp_driver = ogr.GetDriverByName("Memory")
            self._ds = shp_driver.CopyDataSource(data_source, '')
        if self._ds is None:
            raise ValueError('无法获取数据源！')

        return self

    def set_layer(self, layer: ogr.Layer):
        """
        设置图层，正确处理图层复制
        """
        if self._ds is None:
            # 获取内存驱动
            driver = ogr.GetDriverByName('Memory')
            if driver is None:
                raise RuntimeError("Memory 驱动不可用（检查 GDAL 安装）")
            # 创建空的内存 DataSource（类型为 ogr.DataSource，继承 gdal.Dataset）
            self._ds = driver.CreateDataSource('')
            if self._ds is None:
                raise ValueError('无法创建数据源！')
        self._ds.CopyLayer(layer, layer.GetName())
        return self

    @require_attr_not_empty('layer')
    def set_features(self, features: Iterable[ogr.Feature]):
        """
        设置要素features
        :param features:
        :return:
        """
        for feature in features:
            self.layer.CreateFeature(feature)

    @require_attr_not_empty('layer')
    def reproject(self, dst_srs: str | int, src_srs=None) -> 'OpenShp':
        """
        投影转换
        Args:
            dst_srs: 目标空间参考系，可以是EPSG代码(int)、"EPSG:代码"、WKT字符串
            src_srs: 源空间参考系，默认使用shp自带参考系，当shp中没有参考系时，该参数必须设置
        Returns:
            OpenShp对象或None
        """
        # 获取源空间参考系统
        src_sr = src_srs or self.geo_projection
        if src_sr is None:
            raise ValueError("源数据没有空间参考系统，请设置 src_srs 源空间参考系参数！")
        src_sr = GeoConverter.geo_projection2wkt(src_sr)
        self._ds = gdal.VectorTranslate(
            '',  # 输出到内存
            self.data_source,  # 输入数据源
            format="Memory",  # 内存模式
            dstSRS=GeoConverter.geo_projection2wkt(dst_srs),
            srcSRS=src_sr,
        )
        return self

    @require_attr_not_empty('layer')
    def add_other_shp(self, other_shp: 'str|OpenShp'):
        if isinstance(other_shp, str):
            other_shp = OpenShp(other_shp)
        if not other_shp.layer:
            raise ValueError("其他数据源没有图层！")
        for feature in other_shp.layer:
            self.layer.CreateFeature(feature)
        return self

    @require_attr_not_empty('layer')
    def simplify_shp(self, tolerance: float = 0.16):
        for feature in self.layer:  # type: ogr.Feature
            geometry: ogr.Geometry = feature.geometry()
            geometry = geometry.SimplifyPreserveTopology(tolerance)
            feature.SetGeometry(geometry)
            self.layer.SetFeature(feature)
        return self

    def close(self):
        self._ds = None


def open_shp(fp: str, mode: OpenShp.open_mode = 'r') -> OpenShp:
    return OpenShp(fp, mode)
